// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/functions/prestosql/aggregates/arbitrary_aggregate.h>
#include <pollux/exec/aggregate.h>
#include <pollux/expression/function_signature.h>
#include <pollux/functions/lib/aggregates/simple_numeric_aggregate.h>
#include <pollux/functions/lib/aggregates/single_value_accumulator.h>
#include <pollux/functions/prestosql/aggregates/aggregate_names.h>

using namespace kumo::pollux::functions::aggregate;

namespace kumo::pollux::aggregate::prestosql {
    namespace {
        // Arbitrary aggregate returns any arbitrary non-NULL value.
        // We always keep the first (non-NULL) element seen.
        template<typename T>
        class ArbitraryAggregate : public SimpleNumericAggregate<T, T, T> {
            using BaseAggregate = SimpleNumericAggregate<T, T, T>;

        public:
            explicit ArbitraryAggregate(TypePtr resultType) : BaseAggregate(resultType) {
            }

            int32_t accumulatorFixedWidthSize() const override {
                return sizeof(T);
            }

            int32_t accumulatorAlignmentSize() const override {
                return 1;
            }

            void extractValues(char **groups, int32_t numGroups, VectorPtr *result)
            override {
                BaseAggregate::doExtractValues(groups, numGroups, result, [&](char *group) {
                    return *BaseAggregate::Aggregate::template value<T>(group);
                });
            }

            void addRawInput(
                char **groups,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /*unused*/) override {
                DecodedVector decoded(*args[0], rows);

                if (decoded.isConstantMapping()) {
                    auto begin = rows.begin();
                    if (decoded.is_null_at(begin)) {
                        return;
                    }
                    auto value = decoded.value_at<T>(begin);
                    rows.applyToSelected([&](vector_size_t i) {
                        if (exec::Aggregate::isNull(groups[i])) {
                            updateValue(groups[i], value);
                        }
                    });
                } else if (decoded.may_have_nulls()) {
                    rows.applyToSelected([&](vector_size_t i) {
                        if (!decoded.is_null_at(i) && exec::Aggregate::isNull(groups[i])) {
                            updateValue(groups[i], decoded.value_at<T>(i));
                        }
                    });
                } else {
                    rows.applyToSelected([&](vector_size_t i) {
                        if (exec::Aggregate::isNull(groups[i])) {
                            updateValue(groups[i], decoded.value_at<T>(i));
                        }
                    });
                }
            }

            void addIntermediateResults(
                char **groups,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool mayPushdown) override {
                addRawInput(groups, rows, args, mayPushdown);
            }

            void addSingleGroupRawInput(
                char *group,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /*unused*/) override {
                if (!exec::Aggregate::isNull(group)) {
                    return;
                }
                DecodedVector decoded(*args[0], rows);
                auto begin = rows.begin();

                if (decoded.isConstantMapping()) {
                    if (decoded.is_null_at(begin)) {
                        return;
                    }
                    updateValue(group, decoded.value_at<T>(begin));
                } else if (!decoded.may_have_nulls()) {
                    updateValue(group, decoded.value_at<T>(begin));
                } else {
                    // Find the first non-null value.
                    rows.testSelected([&](vector_size_t i) {
                        if (!decoded.is_null_at(i)) {
                            updateValue(group, decoded.value_at<T>(i));
                            return false; // Stop
                        }
                        return true; // Continue
                    });
                }
            }

            void addSingleGroupIntermediateResults(
                char *group,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool mayPushdown) override {
                addSingleGroupRawInput(group, rows, args, mayPushdown);
            }

        protected:
            void initializeNewGroupsInternal(
                char **groups,
                melon::Range<const vector_size_t *> indices) override {
                exec::Aggregate::setAllNulls(groups, indices);
            }

        private:
            inline void updateValue(char *group, T value) {
                exec::Aggregate::clearNull(group);
                *exec::Aggregate::value<T>(group) = value;
            }
        };

        /// Override 'accumulatorAlignmentSize' for UnscaledLongDecimal values as it
        /// uses int128_t type. Some CPUs don't support misaligned access to int128_t
        /// type.
        template<>
        inline int32_t ArbitraryAggregate<int128_t>::accumulatorAlignmentSize() const {
            return static_cast<int32_t>(sizeof(int128_t));
        }

        // Arbitrary for non-numeric types. We always keep the first (non-NULL) element
        // seen. Arbitrary (x) will produce partial and final aggregations of type x.
        class NonNumericArbitrary : public exec::Aggregate {
        public:
            explicit NonNumericArbitrary(const TypePtr &resultType)
                : exec::Aggregate(resultType) {
            }

            // We use singleValueAccumulator to save the results for each group. This
            // struct will allow us to save variable-width value.
            int32_t accumulatorFixedWidthSize() const override {
                return sizeof(SingleValueAccumulator);
            }

            void extractValues(char **groups, int32_t numGroups, VectorPtr *result)
            override {
                POLLUX_CHECK(result);
                (*result)->resize(numGroups);

                auto *rawNulls = exec::Aggregate::getRawNulls(result->get());

                for (int32_t i = 0; i < numGroups; ++i) {
                    char *group = groups[i];
                    auto accumulator = value<SingleValueAccumulator>(group);
                    if (!accumulator->hasValue()) {
                        (*result)->set_null(i, true);
                    } else {
                        exec::Aggregate::clearNull(rawNulls, i);
                        accumulator->read(*result, i);
                    }
                }
            }

            void extractAccumulators(char **groups, int32_t numGroups, VectorPtr *result)
            override {
                extractValues(groups, numGroups, result);
            }

            void addRawInput(
                char **groups,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /*unused*/) override {
                DecodedVector decoded(*args[0], rows, true);
                if (decoded.isConstantMapping() && decoded.is_null_at(rows.begin())) {
                    // nothing to do; all values are nulls
                    return;
                }

                const auto *indices = decoded.indices();
                const auto *baseVector = decoded.base();
                rows.applyToSelected([&](vector_size_t i) {
                    if (decoded.is_null_at(i)) {
                        return;
                    }
                    auto *accumulator = value<SingleValueAccumulator>(groups[i]);
                    if (!accumulator->hasValue()) {
                        accumulator->write(baseVector, indices[i], allocator_);
                    }
                });
            }

            void addIntermediateResults(
                char **groups,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool mayPushdown) override {
                addRawInput(groups, rows, args, mayPushdown);
            }

            void addSingleGroupRawInput(
                char *group,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /*unused*/) override {
                auto *accumulator = value<SingleValueAccumulator>(group);
                if (accumulator->hasValue()) {
                    return;
                }

                DecodedVector decoded(*args[0], rows, true);
                if (decoded.isConstantMapping() && decoded.is_null_at(rows.begin())) {
                    // nothing to do; all values are nulls
                    return;
                }

                const auto *indices = decoded.indices();
                const auto *baseVector = decoded.base();
                // Find the first non-null value.
                rows.testSelected([&](vector_size_t i) {
                    if (!decoded.is_null_at(i)) {
                        accumulator->write(baseVector, indices[i], allocator_);
                        return false; // Stop
                    }
                    return true; // Continue
                });
            }

            void addSingleGroupIntermediateResults(
                char *group,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool mayPushdown) override {
                addSingleGroupRawInput(group, rows, args, mayPushdown);
            }

        protected:
            // Initialize each group, we will not use the null flags because
            // SingleValueAccumulator has its own flag.
            void initializeNewGroupsInternal(
                char **groups,
                melon::Range<const vector_size_t *> indices) override {
                for (auto i: indices) {
                    new(groups[i] + offset_) SingleValueAccumulator();
                }
            }

            void destroyInternal(melon::Range<char **> groups) override {
                for (auto group: groups) {
                    if (isInitialized(group)) {
                        value<SingleValueAccumulator>(group)->destroy(allocator_);
                    }
                }
            }
        };
    } // namespace

    void registerArbitraryAggregate(
        const std::string &prefix,
        bool withCompanionFunctions,
        bool overwrite) {
        std::vector<std::shared_ptr<exec::AggregateFunctionSignature> > signatures{
            exec::AggregateFunctionSignatureBuilder()
            .typeVariable("T")
            .returnType("T")
            .intermediateType("T")
            .argumentType("T")
            .build()
        };

        std::vector<std::string> names = {prefix + kArbitrary, prefix + kAnyValue};
        exec::registerAggregateFunction(
            names,
            std::move(signatures),
            [name = names.front()](
        core::AggregationNode::Step step,
        const std::vector<TypePtr> &argTypes,
        const TypePtr & /*resultType*/,
        const core::QueryConfig & /*config*/)
        -> std::unique_ptr<exec::Aggregate> {
                POLLUX_CHECK_LE(argTypes.size(), 1, "{} takes only one argument", name);
                auto inputType = argTypes[0];
                switch (inputType->kind()) {
                    case TypeKind::BOOLEAN:
                        return std::make_unique<ArbitraryAggregate<bool> >(inputType);
                    case TypeKind::TINYINT:
                        return std::make_unique<ArbitraryAggregate<int8_t> >(inputType);
                    case TypeKind::SMALLINT:
                        return std::make_unique<ArbitraryAggregate<int16_t> >(inputType);
                    case TypeKind::INTEGER:
                        return std::make_unique<ArbitraryAggregate<int32_t> >(inputType);
                    case TypeKind::BIGINT:
                        return std::make_unique<ArbitraryAggregate<int64_t> >(inputType);
                    case TypeKind::HUGEINT:
                        if (inputType->isLongDecimal()) {
                            return std::make_unique<ArbitraryAggregate<int128_t> >(inputType);
                        }
                        POLLUX_NYI();
                    case TypeKind::REAL:
                        return std::make_unique<ArbitraryAggregate<float> >(inputType);
                    case TypeKind::DOUBLE:
                        return std::make_unique<ArbitraryAggregate<double> >(inputType);
                    case TypeKind::TIMESTAMP:
                        return std::make_unique<ArbitraryAggregate<Timestamp> >(inputType);
                    case TypeKind::VARBINARY:
                        [[fallthrough]];
                    case TypeKind::VARCHAR:
                        [[fallthrough]];
                    case TypeKind::ARRAY:
                        [[fallthrough]];
                    case TypeKind::MAP:
                        [[fallthrough]];
                    case TypeKind::ROW:
                        [[fallthrough]];
                    case TypeKind::UNKNOWN:
                        return std::make_unique<NonNumericArbitrary>(inputType);
                    default:
                        POLLUX_FAIL(
                            "Unknown input type for {} aggregation {}",
                            name,
                            inputType->kindName());
                }
            },
            withCompanionFunctions,
            overwrite);
    }
} // namespace kumo::pollux::aggregate::prestosql
